﻿using System;
using System.Collections.Generic;
using System.Linq;
using Algoritmia.DataStructures.Maps;
using Algoritmia.DataStructures.PriorityQueues;

namespace Algoritmia.Problems.Compression
{
    public class HuffmanTreeBuilder : HuffmanTreeBuilder1
    {
        private readonly Func<IEnumerable<HuffmanTree>,IPriorityQueue<HuffmanTree>>  _createPriorityQueue;

        public HuffmanTreeBuilder(Func<string, double, HuffmanTree> createBinaryTree = null,
                                  Func<IEnumerable<HuffmanTree>, IPriorityQueue<HuffmanTree>> createPriorityQueue = null)
            : base(createBinaryTree)
        {
            _createPriorityQueue = createPriorityQueue ??
                (it => new Heap<HuffmanTree>(data: it,comparison: (hfn1, hfn2) => hfn1.Frequency.CompareTo(hfn2.Frequency)));
        }

        public override HuffmanTree BuildTree(IMap<string, double> freq)
        {
            var T = _createPriorityQueue(freq.Select(kv => _createHuffmanTree(kv.Key, kv.Value)));

            while (T.Count > 1)
            {
                var leftTree = T.ExtractOptimal();
                var rightTree = T.ExtractOptimal();
                var newTree = _createHuffmanTree("", leftTree.Frequency + rightTree.Frequency);
                newTree.AppendLeft(leftTree);
                newTree.AppendRight(rightTree);
                T.Add(newTree);
            }
            return T.ExtractOptimal();
        }
    }
}
